import os
import sys
import glob
import argparse
from pathlib import Path
from typing import List, Optional, Tuple
from dataclasses import dataclass
from pprint import pprint
# 全局变量：控制详细日志输出
is_verbose = False

def verbose_print(*args, **kwargs):
    """
    详细日志打印函数，根据is_verbose全局变量决定是否输出
    
    Args:
        *args: 传递给print的位置参数
        **kwargs: 传递给print的关键字参数
    """
    if is_verbose:
        print(*args, **kwargs)

import cv2
from PIL import Image
import numpy as np
import torch
 
# Real-ESRGAN imports
try:
    from basicsr.archs.rrdbnet_arch import RRDBNet
    from realesrgan import RealESRGANer
    from realesrgan.archs.srvgg_arch import SRVGGNetCompact
    _HAS_REAL_ESRGAN = True
except Exception:
    _HAS_REAL_ESRGAN = False


def list_images_sorted(path: str) -> List[str]:
    """
    获取指定路径下的图片文件列表，按文件名排序
    
    Args:
        path: 输入路径，可以是文件或目录
        
    Returns:
        排序后的图片文件路径列表
    """
    if os.path.isdir(path):
        exts = ("*.png", "*.jpg", "*.jpeg", "*.bmp", "*.tif", "*.tiff")
        files = []
        for e in exts:
            files.extend(glob.glob(os.path.join(path, e)))
        return sorted(files)
    elif os.path.isfile(path):
        return [path]
    else:
        return []


def image_read(path: str) -> Optional[np.ndarray]:
    """
    以彩色模式读取图片文件
    
    Args:
        path: 图片文件路径
        
    Returns:
        读取到的图片数组，失败时返回 None
    """
    img = Image.open(path)
    img = np.array(img)
    return img


def image_write(path: str, img: np.ndarray):
    """
    以彩色模式写入图片文件，自动创建目录
    
    Args:
        path: 输出文件路径
        img: 要写入的图片数组
        
    Returns:
        写入成功返回 True，失败返回 False
    """
    
    os.makedirs(os.path.dirname(path), exist_ok=True)
    Image.fromarray(img).save(path)


def to_yuv_float(img_rgb: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    将 BGR 格式的图片转换为 YUV 格式，并归一化到 [0,1] 范围
    
    Args:
        img_rgb: RGB 格式的图片数组，值范围 [0,255]
        
    Returns:
        Y, U, V 三个通道的数组，值范围 [0,1]
    """
    img = img_rgb.astype(np.float32) / 255.0
    yuv = cv2.cvtColor(img, cv2.COLOR_RGB2YUV)
    y, u, v = cv2.split(yuv)
    return y, u, v


def from_yuv_float(y: np.ndarray, u: np.ndarray, v: np.ndarray) -> np.ndarray:
    """
    将 YUV 格式的图片转换回 BGR 格式，并恢复到 [0,255] 范围
    
    Args:
        y: Y 通道数组，值范围 [0,1]
        u: U 通道数组，值范围 [0,1]
        v: V 通道数组，值范围 [0,1]
        
    Returns:
        rgb 格式的图片数组，值范围 [0,255]
    """
    # 确保所有通道的尺寸一致
    target_shape = y.shape
    if u.shape != target_shape:
        if len(u.shape) == 2:
            u = cv2.resize(u, (target_shape[1], target_shape[0]), interpolation=cv2.INTER_LINEAR)
        else:
            u = cv2.resize(u[:, :, 0], (target_shape[1], target_shape[0]), interpolation=cv2.INTER_LINEAR)
    
    if v.shape != target_shape:
        if len(v.shape) == 2:
            v = cv2.resize(v, (target_shape[1], target_shape[0]), interpolation=cv2.INTER_LINEAR)
        else:
            v = cv2.resize(v[:, :, 0], (target_shape[1], target_shape[0]), interpolation=cv2.INTER_LINEAR)
    
    yuv = cv2.merge([y, u, v])
    rgb = cv2.cvtColor(yuv, cv2.COLOR_YUV2RGB)
    rgb = np.clip(rgb, 0.0, 1.0)
    return (rgb * 255.0 + 0.5).astype(np.uint8)


def clamp01(x: float) -> float:
    """
    将数值限制在 [0,1] 范围内
    
    Args:
        x: 输入数值
        
    Returns:
        限制在 [0,1] 范围内的数值
    """
    return float(max(0.0, min(1.0, x)))


class BaseDenoiser:
    """
    降噪器基类，定义了所有降噪器的通用接口
    """
    def __init__(self, strength: float = 0.5):
        """
        初始化降噪器
        
        Args:
            strength: 降噪强度，范围 [0,1]
        """
        self.strength = clamp01(strength)

    def execute(self, input_image: np.ndarray) -> np.ndarray:
        """
        执行降噪处理
        
        Args:
            input_image: 输入图片
            
        Returns:
            降噪后的图片
        """
        raise NotImplementedError


class TemporalDenoiser(BaseDenoiser):
    """
    时域降噪器：通过叠加上一帧数据实现平滑过渡
    """
    def __init__(self, strength: float, ref_image_paths:List[str] = [], window_size = [5,5]):
        """
        初始化时域降噪器
        
        Args:
            strength: 降噪强度，范围 [0,1]
        """
        super().__init__(strength)
        self.history = None
        self.window_size = window_size
        print(f"初始化时域降噪器, strength: {strength}, window_size: {window_size}")
        for i, path in enumerate(ref_image_paths):
            verbose_print(f"加载参考图像：{path}, {i+1}/{len(ref_image_paths)}")
            self.execute(image_read(path))

 
    def _window_sliding(self, image: np.ndarray, ref: np.ndarray):
        """
        优化版本：使用OpenCV形态学操作快速计算邻域范围，保持颜色限制逻辑
        """
        assert image.shape == ref.shape, "image 与 ref 形状必须一致"
        
        # 转为浮点以避免整数截断
        img = image.astype(np.float32, copy=False)
        ref = ref.astype(np.float32, copy=False)
        out = np.empty_like(img)
        
        # 如果降噪强度为0，直接返回当前图像
        if self.strength == 0.0:
            out[:] = img[:]
            return out
        
        # 创建结构元素 - 使用矩形结构元素，大小约为原始窗口的一半
        kernel_size = max(1, min(self.window_size) // 2)
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_size, kernel_size))
        
        if img.ndim == 3 and img.shape[2] > 1:
            # 多通道处理
            num_channels = img.shape[2]
            
            for c in range(num_channels):
                # 计算参考图像的局部最小值（腐蚀）和最大值（膨胀）
                ref_local_min = cv2.erode(ref[:, :, c], kernel, iterations=1)
                ref_local_max = cv2.dilate(ref[:, :, c], kernel, iterations=1)
                
                # 计算当前图像的局部最小值（腐蚀）和最大值（膨胀）
                img_local_min = cv2.erode(img[:, :, c], kernel, iterations=1)
                img_local_max = cv2.dilate(img[:, :, c], kernel, iterations=1)
                
                # 融合两个图像的局部统计信息
                # 取参考图像和当前图像邻域最小值中的更小值
                final_min = np.minimum(ref_local_min, img_local_min)
                # 取参考图像和当前图像邻域最大值中的更大值
                final_max = np.maximum(ref_local_max, img_local_max)
                
                # 应用颜色限制：将参考值限制在计算出的范围内
                ref_channel = ref[:, :, c]
                ref_channel_clipped = np.clip(ref_channel, final_min, final_max)
                
                # 按强度融合当前图像和受限后的参考图像
                out[:, :, c] = img[:, :, c] * (1.0 - self.strength) + ref_channel_clipped * self.strength
                
        else:
            # 单通道处理
            # 计算参考图像的局部最小值（腐蚀）和最大值（膨胀）
            ref_local_min = cv2.erode(ref, kernel, iterations=1)
            ref_local_max = cv2.dilate(ref, kernel, iterations=1)
            
            # 计算当前图像的局部最小值（腐蚀）和最大值（膨胀）
            img_local_min = cv2.erode(img, kernel, iterations=1)
            img_local_max = cv2.dilate(img, kernel, iterations=1)
            
            # 融合两个图像的局部统计信息
            final_min = np.minimum(ref_local_min, img_local_min)
            final_max = np.maximum(ref_local_max, img_local_max)
            
            # 应用颜色限制和强度融合
            ref_clipped = np.clip(ref, final_min, final_max)
            out = img * (1.0 - self.strength) + ref_clipped * self.strength
        
        # 转回原类型（若是整数类型，进行四舍五入与裁剪）
        if np.issubdtype(image.dtype, np.integer):
            info = np.iinfo(image.dtype)
            out = np.clip(np.rint(out), info.min, info.max).astype(image.dtype)
        else:
            out = out.astype(image.dtype, copy=False)
        
        return out
    
    def execute(self, input_image: np.ndarray) -> np.ndarray:
        """
        执行时域降噪
        
        Args:
            input_image: 输入图片（当前帧）
            
        Returns:
            降噪后的图片
        """
        if self.history is None:
            self.history = input_image
        out = self._window_sliding(input_image, self.history)
        self.history = out
        return out


class BilateralDenoiser(BaseDenoiser):
    """
    双边滤波降噪器：结合空间距离和颜色相似性的非线性滤波
    - 在空间域和颜色域同时进行加权平均
    - 能够在平滑噪声的同时保持边缘清晰
    - 适合处理高斯噪声和椒盐噪声
    """
    def __init__(self, strength: float = 0.5, sigma_space: Optional[float] = 45, sigma_color: Optional[float] = 45, diameter: int = 3):
        """
        初始化双边滤波降噪器
        
        Args:
            strength: 降噪强度，范围 [0,1]
            sigma_space: 空间域标准差，如果为 None 则根据强度自动计算
            sigma_color: 颜色域标准差，如果为 None 则根据强度自动计算
            diameter: 滤波器直径，0 表示自动计算
        """
        super().__init__(strength)
        verbose_print(f"初始化双边滤波降噪器, strength: {strength}, sigma_space: {sigma_space}, sigma_color: {sigma_color}, diameter: {diameter}")
        # 合理范围映射
        self.sigma_space = 3 + self.strength * 10 if sigma_space is None else float(sigma_space)
        self.sigma_color = 20 + self.strength * 100 if sigma_color is None else float(sigma_color)
        self.diameter = int(diameter)

    def execute(self, input_image: np.ndarray) -> np.ndarray:
        """
        执行双边滤波降噪
        
        Args:
            input_image: 输入图片
            
        Returns:
            降噪后的图片
        """
        verbose_print("开始双边滤波降噪...")
        img = input_image
        img32 = img.astype(np.float32)
        out = cv2.bilateralFilter(img32, d=self.diameter, sigmaColor=self.sigma_color, sigmaSpace=self.sigma_space, borderType=cv2.BORDER_REFLECT)
        out = np.clip(out, 0, 255).astype(np.uint8)
        return out


class GuidedFilterDenoiser(BaseDenoiser):
    """
    引导滤波降噪器：使用引导图像进行边缘保持滤波
    - 利用引导图像的结构信息来指导滤波过程
    - 在平滑区域进行强滤波，在边缘区域保持细节
    - 计算效率高，适合实时处理
    """
    def __init__(self, strength: float = 0.5, radius: Optional[int] = None, eps: Optional[float] = None):
        """
        初始化引导滤波降噪器
        
        Args:
            strength: 降噪强度，范围 [0,1]
            radius: 滤波窗口半径，如果为 None 则根据强度自动计算
            eps: 正则化参数，如果为 None 则根据强度自动计算
        """
        super().__init__(strength)
        verbose_print(f"初始化引导滤波降噪器, strength: {self.strength}, radius: {self.radius}, eps: {self.eps}")
        # 半径与 eps 随强度映射
        self.radius = int(round(2 + self.strength * 6)) if radius is None else int(radius)
        # eps 对应 8-bit 归一化范围，这里用 1e-6..1e-2 的区间
        self.eps = (1e-4 + self.strength * (5e-3)) if eps is None else float(eps)
        # 检测 ximgproc
        self._has_ximgproc = hasattr(cv2, 'ximgproc') and hasattr(cv2.ximgproc, 'guidedFilter')

    def execute(self, input_image: np.ndarray) -> np.ndarray:
        """
        执行引导滤波降噪
        
        Args:
            input_image: 输入图片
            
        Returns:
            降噪后的图片
        """
        verbose_print("开始引导滤波降噪...")
        img = input_image
        if self._has_ximgproc:
            guide = img
            src = img
            # OpenCV guidedFilter expects CV_8U or CV_32F
            guide32 = guide.astype(np.float32) / 255.0
            src32 = src.astype(np.float32) / 255.0
            out32 = cv2.ximgproc.guidedFilter(guide32, src32, self.radius, self.eps)
            out = np.clip(out32 * 255.0 + 0.5, 0, 255).astype(np.uint8)
            return out
        else:
            # 回退：使用 edgePreservingFilter 作为近似
            out = cv2.edgePreservingFilter(img, flags=1, sigma_s=max(10, int(self.radius * 4)), sigma_r=min(0.4, 0.02 + self.strength * 0.2))
            return out


class HistEQDenoiser(BaseDenoiser):
    """
    直方图均衡化降噪器：通过增强对比度来改善图像质量
    - 使用 CLAHE（对比度受限的自适应直方图均衡化）
    - 在局部区域内进行直方图均衡化，避免全局过度增强
    - 特别适合低光照条件下的图像增强
    """
    def __init__(self, strength: float = 0.5, clip_limit: Optional[float] = 1.0, tile_grid_size: Optional[Tuple[int, int]] = (8,8)):
        """
        初始化直方图均衡化降噪器
        
        Args:
            strength: 降噪强度，范围 [0,1]
            clip_limit: 对比度限制，如果为 None 则根据强度自动计算
            tile_grid_size: 网格大小，如果为 None 则根据强度自动计算
        """
        super().__init__(strength)
        self.clip_limit = clip_limit
        self.tile_grid_size = tile_grid_size
        verbose_print(f"初始化直方图均衡化降噪器, strength: {self.strength}, clip_limit: {clip_limit}, tile_grid_size: {tile_grid_size}")
        # 使用 CLAHE（更鲁棒），强度映射到 clipLimit


    def execute(self, input_image: np.ndarray) -> np.ndarray:
        """
        执行直方图均衡化降噪
        
        Args:
            input_image: 输入图片
            
        Returns:
            增强后的图片
        """
        verbose_print("开始直方图均衡化...")
        img = input_image
        y, u, v = to_yuv_float(img)
        y_u8 = np.clip(y * 255.0 + 0.5, 0, 255).astype(np.uint8)
        clahe = cv2.createCLAHE(clipLimit=self.clip_limit, tileGridSize=self.tile_grid_size)
        y_eq = clahe.apply(y_u8).astype(np.float32) / 255.0
        out = from_yuv_float(y_eq, u, v)
        return out


class RealESRGANDenoiser(BaseDenoiser):
    """
    Real-ESRGAN 降噪器：基于深度学习的超分辨率和降噪
    - 使用 RRDBNet 网络进行图像增强
    - 支持单帧处理，无需多帧信息
    - 能够同时进行降噪、去模糊和轻微超分辨率
    """
    def __init__(
        self,
        strength: float = 0.5,
        model_name: str = 'RealESRGAN_x4plus',
        model_path: Optional[str] = None,
        device: str = 'cuda',
        tile: int = 0,
        tile_pad: int = 10,
        pre_pad: int = 0,
        fp32: bool = False
    ):
        """
        初始化 Real-ESRGAN 降噪器
        
        Args:
            strength: 降噪强度，范围 [0,1]
            model_name: 模型名称
            model_path: 模型文件路径，如果为 None 则使用默认路径
            device: 计算设备 ('cuda' 或 'cpu')
            tile: 瓦片大小，0 表示不使用瓦片
            tile_pad: 瓦片填充
            pre_pad: 预填充
            fp32: 是否使用 FP32 精度
        """
        super().__init__(strength)
        self.enabled = _HAS_REAL_ESRGAN
        self.upsampler = None
        if not self.enabled:
            return

        self.model_name = model_name
        self.device = device
        self.tile = int(tile)
        self.tile_pad = int(tile_pad)
        self.pre_pad = int(pre_pad)
        self.fp32 = bool(fp32)

        # 根据名字创建模型与默认路径
        model = None
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
        default_path = "D:/source/SR/Real-ESRGAN/weights/RealESRGAN_x4plus.pth"
        mpath = model_path if model_path else default_path
        try:
            self.upsampler = RealESRGANer(
                scale=1,
                model_path=mpath,
                model=model,
                tile=self.tile,
                tile_pad=self.tile_pad,
                pre_pad=self.pre_pad,
                half=not self.fp32,
                device=self.device
            )
        except Exception as e:
            verbose_print(f"[Warn] RealESRGAN init failed: {e}", file=sys.stderr)
            self.upsampler = None

    def execute(self, input_image: np.ndarray) -> np.ndarray:
        """
        执行 Real-ESRGAN 降噪
        
        Args:
            input_image: 输入图片
            
        Returns:
            增强后的图片
        """
        verbose_print("开始 Real-ESRGAN 降噪...")
        img = input_image
        if not self.enabled or self.upsampler is None:
            # 直接返回，作为"无操作"降级
            return img.copy()
        try:
            # outscale 使用 1 表示仅去噪/增强，不放大
            out, _ = self.upsampler.enhance(img, outscale=1)
            return out
        except Exception as e:
            verbose_print(f"[Warn] RealESRGAN enhance failed: {e}", file=sys.stderr)
            return img.copy()


class DenoiseManager:
    """
    降噪管理器：协调整个降噪流程
    - 管理输入输出路径
    - 构建降噪器管道
    - 处理多帧图像序列
    - 支持详细的运行日志
    """
    def __init__(self, input_path: str, ref_images: List[str], output_path: str, strength: float, methods: List[str]):
        """
        初始化降噪管理器
        
        Args:
            input_path: 输入路径
            ref_images: 参考图片path或dir
            output_path: 输出路径
            strength: 降噪强度
            methods: 降噪器方法列表
        """
        verbose_print(f'''初始化降噪管理器
                      input_path: {input_path}, 
                      ref_images: {ref_images}, 
                      output_path: {output_path}, 
                      strength: {strength}, 
                      methods: {methods}''')
        self.input_path = input_path
        self.ref_image_paths = self.get_all_files(ref_images, ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif', '.webp'])
        self.output_path = output_path
        self.strength = clamp01(strength)
        self.methods = methods
        # 支持的图片格式
        self.pipeline: List[BaseDenoiser] = self._build_pipeline(methods)

    def get_all_files(self, paths, vaild_suffixs=[]):
        """检查文件是否为支持的图片格式"""
        files = []
        if paths is not None:
            for path in paths:
                if os.path.isfile(path):
                    files.append(path)
                else:
                    for root, dirs, filenames in os.walk(path):
                        for filename in filenames:
                            filepath = os.path.join(root, filename)
                            suffix = os.path.splitext(filename)[1].lower()
                            if suffix in vaild_suffixs:
                                files.append(filepath)
        return files

    def _build_pipeline(self, methods: List[str]) -> List[BaseDenoiser]:
        """
        构建降噪器管道
        
        Args:
            methods: 降噪器方法列表
            
        Returns:
            降噪器管道列表
        """
        # 方法名不区分大小写，支持：temporal, bilateral, guided, histeq, realesrgan
        pipe: List[BaseDenoiser] = []
        for m in methods:
            key = m.strip().lower()
            if key == 'temporal':
                denoiser = TemporalDenoiser(strength=self.strength, ref_image_paths = self.ref_image_paths)
                pipe.append(denoiser)
                verbose_print(f"[Verbose] 添加时域降噪器，强度: {self.strength}, 参考图片: {len(self.ref_image_paths)}")
            elif key == 'bilateral':
                denoiser = BilateralDenoiser(strength=self.strength)
                pipe.append(denoiser)
                verbose_print(f"[Verbose] 添加双边滤波降噪器，强度: {self.strength}, sigma_space: {denoiser.sigma_space}, sigma_color: {denoiser.sigma_color}")
            elif key == 'guided':
                denoiser = GuidedFilterDenoiser(strength=self.strength)
                pipe.append(denoiser)
                verbose_print(f"[Verbose] 添加引导滤波降噪器，强度: {self.strength}, 半径: {denoiser.radius}, eps: {denoiser.eps}")
            elif key in ('histeq', 'clahe'):
                denoiser = HistEQDenoiser(strength=self.strength)
                pipe.append(denoiser)
                verbose_print(f"[Verbose] 添加直方图均衡化降噪器，强度: {self.strength}, clip_limit: {denoiser.clip_limit}, tile_size: {denoiser.tile_grid_size}")
            elif key == 'realesrgan':
                denoiser = RealESRGANDenoiser(strength=self.strength)
                pipe.append(denoiser)
                status = "已启用" if denoiser.enabled and denoiser.upsampler is not None else "未启用"
                verbose_print(f"[Verbose] 添加Real-ESRGAN降噪器，强度: {self.strength}, 状态: {status}")
            elif key == '' or key == 'none':
                # 跳过空名称
                continue
            else:
                verbose_print(f"[Warn] Unknown method '{m}', skipped.", file=sys.stderr)
        return pipe

    def _infer_output_path(self, output_path: str, default_name = 'denoise.png'):
        """
        推断输出文件路径
        
        Args:
            ipath: 输入文件路径
            
        Returns:
            输出文件路径
        """
        path = Path(output_path)
        if path.suffix:  # 有文件扩展名 → 文件路径
            parent = path.parent
            if not parent.exists():
                os.makedirs(parent)
        else:            # 无文件扩展名 → 目录路径
            if not path.exists():
                os.makedirs(path)
            path = path / default_name
        return str(path)



    def run(self):
        """
        运行降噪流程
        """
        img = image_read(self.input_path)
        if img is None:
            raise ValueError(f"Failed to read image from '{self.input_path}'")
        out_image = img.copy()
        for i, denoiser in enumerate(self.pipeline):
            denoiser_name = type(denoiser).__name__
            verbose_print(f"[Verbose] 应用 {denoiser_name} 降噪器...")
            # 对每个降噪器传入当前图像
            out_image = denoiser.execute(out_image)
            verbose_print(f"[Verbose] {denoiser_name} 降噪完成")

        out_path = self._infer_output_path(self.output_path)
        print(f"Output saved to '{out_path}'.")
        image_write(out_path, out_image)


def parse_args():
    """
    解析命令行参数
    
    Returns:
        解析后的参数对象
    """
    parser = argparse.ArgumentParser(description="Modular image/video denoiser with combinable methods.")
    parser.add_argument("input", type=str, help="输入图片文件或目录")
    parser.add_argument("-r", "--ref_images", type=str, nargs='+', help="参考图片或图片目录")
    parser.add_argument("-o", "--output", type=str, default='result.png', help="输出图片或目录")
    parser.add_argument("-s", "--strength", type=float, default=0.8, help="降噪强度 [0,1]")
    parser.add_argument("-m", "--methods", type=str, nargs="+", default=["temporal", "bilateral"] , choices=["temporal", "bilateral", "guided", "histeq", "realesrgan"], help="降噪算法名称, 默认[temporal、bilateral、histeq]")
    parser.add_argument("-v", "--verbose", action="store_true", help="打印详细运行日志")
    return parser.parse_args()


def main():
    """
    主函数
    """
    args = parse_args()
    # 设置全局详细日志变量
    global is_verbose
    is_verbose = args.verbose
    
    mgr = DenoiseManager(
        input_path=args.input,
        ref_images=args.ref_images,
        output_path=args.output,
        strength=args.strength,
        methods=args.methods
    )
    with torch.no_grad():
        mgr.run()
    # 清理显存
    if torch.cuda.is_available():
        torch.cuda.synchronize()
        torch.cuda.empty_cache()


if __name__ == "__main__":
    main()
